import numpy as np
import os
import ntpath
import time

import torch
from . import util
from . import html

try:
    from torch.utils.tensorboard import SummaryWriter
except ImportError:
    print("Warning: tensorboard library not found. TensorBoard logging will be disabled.")
    print("Please install it with: pip install tensorboard")
    SummaryWriter = None


class Visualizer():
    def __init__(self, opt):
        self.tf_log = opt.tf_log
        self.use_html = opt.isTrain and not opt.no_html
        self.win_size = opt.display_winsize
        self.name = opt.name
        self.writer = None

        if SummaryWriter is not None:
            self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, 'logs')
            self.writer = SummaryWriter(log_dir=self.log_dir)

        self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web')
        self.img_dir = os.path.join(self.web_dir, 'images')

        self.log_name = os.path.join(opt.checkpoints_dir, opt.name, 'loss_log.txt')
        with open(self.log_name, "a") as log_file:
            now = time.strftime("%c")
            log_file.write('================ Training Loss (%s) ================\n' % now)

    # |visuals|: dictionary of images to display or save
    def display_current_results(self, visuals, epoch, step, types=None):
        if types is None:
            types = ['tf', 'web']
        output_types = set(types)

        if 'tf' in output_types and self.writer:
            for label, image_numpy in visuals.items():
                image_tensor = None
                if image_numpy.ndim == 2:  # Grayscale HxW
                    # Convert HxW uint8 numpy to 1xHxW uint8 Tensor
                    image_tensor = torch.from_numpy(image_numpy).unsqueeze(0)
                elif image_numpy.ndim == 3:  # Color HxWxChannels
                    # Convert HxWxChannels uint8 numpy to ChannelsxHxW uint8 Tensor
                    image_tensor = torch.from_numpy(image_numpy).permute(2, 0, 1)
                else:
                    print(f"Warning: Skipping image '{label}' "
                          f"with unexpected dimensions for TensorBoard: {image_numpy.shape}")
                    continue  # Skip if dimensions are not 2 or 3
                if image_tensor.dtype != torch.uint8:
                    image_tensor = image_tensor.to(torch.uint8)

                self.writer.add_image(f"Samples/{label}", img_tensor=image_tensor, global_step=step)
            self.writer.flush()

        if 'web' in output_types:
            util.mkdirs([self.web_dir, self.img_dir])
            for label, image_numpy in visuals.items():
                if isinstance(image_numpy, list):
                    for i in range(len(image_numpy)):
                        img_path = os.path.join(self.img_dir, 'epoch%.3d_%s_%d.jpg' % (epoch, label, i))
                        util.save_image(image_numpy[i], img_path)
                else:
                    img_path = os.path.join(self.img_dir, 'epoch%.3d_%s.jpg' % (epoch, label))
                    util.save_image(image_numpy, img_path)

            # update website
            webpage = html.HTML(self.web_dir, 'Experiment name = %s' % self.name, refresh=30)
            for n in range(epoch, 0, -1):
                webpage.add_header('epoch [%d]' % n)
                ims = []
                txts = []
                links = []

                for label, image_numpy in visuals.items():
                    if isinstance(image_numpy, list):
                        for i in range(len(image_numpy)):
                            img_path = 'epoch%.3d_%s_%d.jpg' % (n, label, i)
                            ims.append(img_path)
                            txts.append(label + str(i))
                            links.append(img_path)
                    else:
                        img_path = 'epoch%.3d_%s.jpg' % (n, label)
                        ims.append(img_path)
                        txts.append(label)
                        links.append(img_path)
                if len(ims) < 10:
                    webpage.add_images(ims, txts, links, width=self.win_size)
                else:
                    num = int(round(len(ims) / 2.0))
                    webpage.add_images(ims[:num], txts[:num], links[:num], width=self.win_size)
                    webpage.add_images(ims[num:], txts[num:], links[num:], width=self.win_size)
            webpage.save()

    def log(self, epoch, i, errors, t, step, types=None):
        """
        将当前误差记录或绘制到指定的目标位置。

        Args:
            epoch (int): 当前的 epoch 数。
            i (int): 当前 epoch 内的迭代次数。
            errors (dict): 误差标签及其对应的数值标量值字典。
                           例如: {'G_GAN': 0.1, 'D_loss': 0.05}
            t (float): 最近 print_freq 次迭代花费的时间 (原始持续时间)。
            step (int): 当前累积的总步数 (例如，处理过的总图像数)。
            types (list): 指定输出目标的字符串列表。
                          可包含 'file'（文本文件）、'console'（控制台）、'tf'（TensorBoard）。
                          默认为 ['tf']。
        """
        if types is None:
            types = ['tf']
        output_types = set(types)

        if 'console' in output_types or 'file' in output_types:
            message = '(epoch: %d, iters: %d, time: %.3f) ' % (epoch, i, t)
            for k, v in errors.items():
                if v != 0:
                    message += '%s: %.3f ' % (k, v)

            if 'console' in output_types:
                print(message)

            if 'file' in output_types:
                with open(self.log_name, "a") as log_file:
                    log_file.write('%s\n' % message)

        if 'tf' in output_types:
            for tag, value in errors.items():
                scalar_value = None
                # 确保值是标量数字，然后添加到 TensorBoard
                if isinstance(value, (int, float, np.number)):
                    scalar_value = value
                elif isinstance(value, torch.Tensor) and value.numel() == 1:
                    scalar_value = value.item()

                if scalar_value is not None:
                    self.writer.add_scalar(f"Loss/{tag}", scalar_value=scalar_value, global_step=step)

    # save image to the disk
    def save_images(self, webpage, visuals, image_path):
        image_dir = webpage.get_image_dir()
        short_path = ntpath.basename(image_path[0])
        name = os.path.splitext(short_path)[0]

        webpage.add_header(name)
        ims = []
        txts = []
        links = []

        for label, image_numpy in visuals.items():
            image_name = '%s_%s.png' % (name, label)
            save_path = os.path.join(image_dir, image_name)
            util.save_image(image_numpy, save_path)

            ims.append(image_name)
            txts.append(label)
            links.append(image_name)
        webpage.add_images(ims, txts, links, width=self.win_size)

    def log_validation_metrics(self, epoch, metrics, global_step):
        # Assuming you use TensorBoardX or SummaryWriter

        for k, v in metrics.items():
            self.writer.add_scalar(f'val_metrics/{k}', v, global_step=global_step)
            # You might also want to log validation images
            # self.writer.add_image(...)
        # Or if you have other logging mechanisms (e.g., custom logging)
        # print(f"Validation Epoch {epoch}: {metrics}")
